import os
from contextlib import contextmanager

from typing import Any, Optional, Callable, Iterator, Dict


def get_parameter_from_env(
    k: str, default: Any, type_converter: Optional[Callable] = None
) -> Any:
    """Get parameter value from system environment variable

    Parameters
    ----------
    k : str
        System environment variable key used to host the parameter.
    default : Any
        Default parameter value if the key does not present in the system environment variables.
    type_converter : Optional[Callable]
        Function used to convert a str value to the desired data type for the parameter,
        default to None, not to convert and return the str value.

    Returns
    -------
    Any
        Parameter value.
    """
    if k in os.environ:
        if type_converter is not None:
            return type_converter(os.environ[k])
        else:
            return os.environ[k]
    else:
        return default


@contextmanager
def inject_env_vars(**kwargs: Any) -> Iterator[None]:
    """A context manager which adds environment variables before running other code,
       and then changes those environment variables back to previous values on exit.

    Parameters
    ----------
    **kwargs : Dict[str, Any]
        Environment variables to set in the form of key, value pairs,
        It is expected that those values can be converted to string using the `str` function.

    Examples
    --------
    >>> with inject_env_vars(SW_DEHAZE_MODEL_VERSION="v2"):
            # call a function which eventually call image_dehaze
            # pass the instruction of using model v2 via environment variable
            image_decloud(tile_id, date_range,result_dir, bands)
        # call again but keep using the default v1 dehaze model
        image_decloud(tile_id, date_range,result_dir, bands)
    """
    prev_env: Dict[str, Optional[str]] = {}
    for k, v in kwargs.items():
        if k in os.environ:
            prev_env[k] = os.environ[k]
        else:
            prev_env[k] = None
        os.environ[k] = str(v)

    try:
        yield
    finally:
        for k, v in prev_env.items():
            if v is None:
                del os.environ[k]
            else:
                os.environ[k] = str(v)
